import numpy as np
import matplotlib.pyplot as plt
#from scipy.optimize import least_squares
from scipy.integrate import quad
from scipy.optimize import curve_fit

T = np.array([300,400,450,500,550,600,680])
D = np.array([0.7,1.3,2.0,2.8,4.2,5.0,7.4])*1e-6
err = np.array([0.2,0.4,1.0,0.4,0.9,2.1,2.3])*1e-6

def func(x, a, b):
    return a*x + b

popt, pcov = curve_fit(func, 1/T, np.log(D), sigma=np.log(D+err)-np.log(D), absolute_sigma=True)
slope = popt[0] 
intercept = popt[1]

slope_0, intercept_0 = np.polyfit(1/T,np.log10(D),deg=1) # plotting
x = 1/np.linspace(200,1000,1000)
y = x*slope_0 + intercept_0

T1 = np.array([200,300,400,500,600])
D1 = np.array([0.0522,0.558,3.52,7.93,15.9])*1e-6
err1 = np.array([2.355550124460,10.16454123021,20.34220609566,119.4635129259,185.9465085803,])*1e-11

popt, pcov = curve_fit(func, 1/T1, np.log(D1), sigma=np.log(D1+err1)-np.log(D1), absolute_sigma=True)
slope1 = popt[0]
intercept1 = popt[1]
#slope1, intercept1 = np.polyfit(1/T1,np.log(D1),deg=1)

slope_1, intercept_1 = np.polyfit(1/T1,np.log10(D1),deg=1) # plotting
x1 = 1/np.linspace(200,1000,1000)
y1 = x1*slope_1 + intercept_1

fig, ax1 = plt.subplots(figsize=(8,6.2))

ax1.errorbar(1000/T,np.log10(D),yerr=np.log10(D+err)-np.log10(D),fmt='o',markersize=8,capsize=5,zorder=2,color='C0',label='BASIS')
#ax1.plot(1000/T,np.log10(D),'o',markersize=8,zorder=2,color='C0',label='BASIS')
ax1.plot(x*1e3,y,color='C0',ls='dashed',zorder=0)

ax1.errorbar(1000/T1,np.log10(D1),yerr=np.log10(D1+err1)-np.log10(D1),fmt='s',markersize=8,capsize=5,zorder=2,color='C1',label='unconstrained')
#ax1.plot(1000/T1,np.log10(D1),'s',markersize=8,zorder=2, color='C1',label='_nolegend_')
ax1.plot(x1*1e3,y1,color='C1',ls='dashed',zorder=1,label='_nolegend_')
#ax1.plot(1000/600,np.log10(7.52*10**(-6)),color='Purple',marker='P',markersize=15,zorder=1)
ax1.plot(1000/600,np.log10(15.9e-6),'s',color='C1',markersize=8,label='_nolegend_',zorder=8)
ax1.plot(1000/600,np.log10(1.60e-5),'d',color='C2',markersize=8,label=r'rigid PS$_4^{3-}$',zorder=10)
ax1.plot(1000/600,np.log10(1.72e-5),'^',color='C5',markersize=8,label='fixed PS$_4^{3-}$ rotation',zorder=10)
ax1.plot(1000/600,np.log10(1.46e-5),'H',color='pink',markersize=8,label=r'fixed free S$^2-$ and Cl$^-$',zorder=9)
ax1.plot(1000/600,np.log10(1.06e-5),'p',color='blue',markersize=8,label=r'fixed PS$_4^{3-}$',zorder=10)
ax1.plot(1000/600,np.log10(0.60e-5),'X',color='black',markersize=8,label='fixed host',zorder=10)

ax1.plot(10/3,np.log10(5.58e-7),'s',color='C1',markersize=8,label='_nolegend_',zorder=8)
ax1.plot(10/3,np.log10(4.62e-7),'d',color='C2',markersize=8,label='_nolegend_',zorder=10)
ax1.plot(10/3,np.log10(3.54e-7),'^',color='C5',markersize=8,label='_nolegend_',zorder=10)
ax1.plot(10/3,np.log10(2.78e-7),'H',color='pink',markersize=8,label='_nolegend_',zorder=9)
ax1.plot(10/3,np.log10(9.94e-8),'p',color='blue',markersize=8,label='_nolegend_',zorder=10)
ax1.plot(10/3,np.log10(4.59e-8),'X',color='black',markersize=8,label='_nolegend_',zorder=10)

ax1.set_xlim(1,4)
ax1.set_xticks([1,2,3,4])
ax1.set_ylim(-8,-4)
ax1.set_yticks([-8,-7,-6,-5,-4])
ax1.set_xlabel('1000/T (1/K)',fontsize=20)
ax1.set_ylabel(r'log10(D) (log10(cm$^2$/s))',fontsize=20)
ax1.tick_params(which='both',direction='in',labelsize=20,pad=8)
ax1.text(2.2,-4.5,r'$E_a (BASIS)$ = %.2f eV' %(slope*(-8.617*1e-2/1000)),fontsize=20,color='C0')
ax1.text(2.2,-4.9,r'$E_a (MLMD)$ = %.2f eV' %(slope1*(-8.617*1e-2/1000)),fontsize=20,color='C1')
#ax1.text(1.2,-5,r'D$_{DOS}$',fontsize=20,color='purple')
ax1.minorticks_on()


ax1.legend(fontsize=14,markerscale=1,handletextpad=0.1,frameon=False,loc='lower left')


ax2 = ax1.twiny()
ax2.set_xlim(1,4)
ax2.set_xlabel('T (K)',fontsize=20,labelpad=-10)
ax2.set_xticks([1000/600,1000/300])
ax2.set_xticklabels([600,300])
ax2.tick_params(which='both',direction='in',labelsize=20)

plt.tight_layout()
plt.savefig('Activation_Energy_v10.pdf',format='pdf',bbox_inches='tight')
plt.show()
